#!/usr/bin/env python
import os,re,urllib,time,zlib,stat
from models import *
import settings

def hardcheck(dir):
    '''
    check the directory is exist. Otherwise make it.
    '''
    if os.path.exists(dir):
        if not os.path.isdir(dir):
            os.remove(dir)
            os.mkdir(dir)
    else:
        os.mkdir(dir)
        
def parsecmd(url1):
    '''
    return False if url says stop, try 10 times, then return True
    '''
    if url1=='':
        return True
    for t1 in range(10):
        try:
            f=urllib.urlopen(url1)
            cont=f.readline()
            f.close()
            if cont.strip().lower()=='stop':
                return False
            else:
                return True
        except:
            time.sleep(1)
    return True

def locktable(tablename):
    ''' lock table for write '''
    metadata.bind.execute('LOCK TABLES %s write'%tablename)

def unlock():
    ''' unlock all tables'''
    metadata.bind.execute('UNLOCK TABLES')

def pdb2pdbqt(PDB):
    ''' use AutoDockTool to prepare pdbqt file'''
    os.system('%s %s -l %s'%(settings.pythonsh,settings.prepare_ligand4,PDB))

def autodock(PDBQT,DLGNAME):
    PRONAME=settings.ProteinName
    pdbqtMdpf(PDBQT,'%s.dpf.temp'%PRONAME,'%s.dpf'%PRONAME)
    os.system('%s -p %s.dpf -l %s'%(autodockcmd(),PRONAME,DLGNAME))

def pdbqtMdpf(PDBQT,DPF,NEWDPF=''):
    '''
    usage:pdbqtMpdf(<pdbqt file>,<dpf file>,<new dpf file>)
    Update dpf file using ligand parameter from pdbqt file.
    ==================
    ligand_types
    move
    about
    ndihe
    torsdof
    ==================
    '''
    # read PDBQT
    fl=open(PDBQT).readlines()
    NDIHE=int(fl[0].split()[1]) # ndihe
    LIGAND_TYPES_DIC={} # ligand_types
    atomnum=0
    sumX=sumY=sumZ=0
    for line in fl[1:]:
        if not line.startswith('ATOM'):
            continue
        linelst=line.split()
        sumX+=float(linelst[5])
        sumY+=float(linelst[6])
        sumZ+=float(linelst[7])
        type1=linelst[11]
        LIGAND_TYPES_DIC[type1]=1
        atomnum+=1
    cenX=sumX/atomnum
    cenY=sumY/atomnum
    cenZ=sumZ/atomnum
    LIGAND_TYPES_STR=''
    for x in LIGAND_TYPES_DIC.keys():
        LIGAND_TYPES_STR+=(x+' ')
    # write DPF
    if NEWDPF=='':
        NEWDPF=DPF
    fDPF=open(DPF)
    fl=fDPF.readlines()
    fDPF.close()
    for lineseq in range(len(fl)):
        line=fl[lineseq]
        linelst=line.split()
        if linelst[0]=='ligand_types':
            fl[lineseq]='ligand_types %s\n'%LIGAND_TYPES_STR
        elif linelst[0]=='move':
            fl[lineseq]='move %s\n'%PDBQT
        elif linelst[0]=='about':
            fl[lineseq]='about %.4f %.4f %.4f\n'%(cenX,cenY,cenZ)
        elif linelst[0]=='ndihe':
            fl[lineseq]='ndihe %d\n'%NDIHE
        elif linelst[0]=='torsdof':
            fl[lineseq]='torsdof %d %s\n'%(NDIHE,linelst[2])
            break # the last midification is 'torsdof'
    fNEWDPF=open(NEWDPF,'w')
    fNEWDPF.writelines(fl)
    fNEWDPF.close()
    
def autodockcmd():
    WORKDIRROOT=settings.TMPDIR
    AUTODOCKCMD=os.path.join(WORKDIRROOT,'autodock')
    if not os.path.exists(AUTODOCKCMD):
        autodock_src,type1=settings.SrcAutoDock
        if type1=='gz':
            tf1=AUTODOCKCMD+'.gz'
            urllib.urlretrieve(autodock_src,tf1)
            os.system(settings.gunzip+' '+tf1)
        elif type1=='bz':
            tf1=AUTODOCKCMD+'.bz2'
            urllib.urlretrieve(autodock_src,tf1)
            os.system(settings.bunzip+' '+tf1)
        else:
            tf1=AUTODOCKCMD
            urllib.urlretrieve(autodock_src,tf1)
        os.chmod(AUTODOCKCMD,stat.S_IXUSR+stat.S_IRUSR+stat.S_IWUSR)
        time.sleep(1)
    return AUTODOCKCMD

def dlg2pdb(DLG,PDBQT,mol1):
    '''
    usage:dlg2pdb(dlg file,template pdbqt)
    Extract pdb coordinates from dlg file.
    The output pdb:
    <name>_<cluster rank>_tnh.pdb
    For Example:
    ZINC00000011.dlg => zinc00000011_1_tnh.pdb, zinc00000011_2_tnh.pdb... 
    They are pdb files without H.    
    Save all the poses and energies into the database.
    '''
    fPDBQT=open(PDBQT)
    fAll=fPDBQT.read()
    fPDBQT.close()
    PDBQTLines=re.findall(r'ATOM .+?\n',fAll,re.S)

    # clean PDB move the atom names from [12:15] to [13:16]
    for lineseq1 in range(len(PDBQTLines)):
        line1=PDBQTLines[lineseq1]
        if not line1.startswith('ATOM'):
            continue
        if line1[12]!=' ' and line1[16]==' ':
            templine=line1[:12]+' '+line1[12:15]+line1[16:]
            PDBQTLines[lineseq1]=templine        
    # end clean
    
    #
    fDlG=open(DLG)
    fAll=fDlG.read()
    fDlG.close()
    # LIGNAME
    LIGNAME=settings.LIGNAME
    # fetch information section
    infoAll=re.search(r'Keeping original residue number \(specified in the input PDBQ file\) for outputting\.\n.+ENDMDL\n\n\n',fAll,re.S).group()
    # split each pose
    for pose_content1 in re.findall(r'MODEL .+?ENDMDL\n',infoAll,re.S):
        # grasp Rank
        CLUSTERRANK=int(re.search(r'USER    Cluster Rank = (?P<pp>\d+)\n',pose_content1).group('pp'))
        # grasp Energy terms of AutoDock
        #USER    Estimated Free Energy of Binding    =   -7.64 kcal/mol  [=(1)+(2)+(3)-(4)]
        #USER    Estimated Inhibition Constant, Ki   =    2.53 uM (micromolar)  [Temperature = 298.15 K]
        #USER    
        #USER    (1) Final Intermolecular Energy     =   -9.32 kcal/mol
        #USER        vdW + Hbond + desolv Energy     =   -9.88 kcal/mol
        #USER        Electrostatic Energy            =   +0.56 kcal/mol
        #USER    (2) Final Total Internal Energy     =   -0.22 kcal/mol
        #USER    (3) Torsional Free Energy           =   +1.65 kcal/mol
        #USER    (4) Unbound System's Energy         =   -0.25 kcal/mol
        feb=float(re.search(r'USER    Estimated Free Energy of Binding    = (?P<pp>[ 0-9.\-+]+) kcal/mol',pose_content1).group('pp'))
        ic,unit=re.search(r'USER    Estimated Inhibition Constant\, Ki   = (?P<pp>[ 0-9.\-+]+) (?P<qq>[num])M',pose_content1).groups()
        if unit=='n':
            ic=float(ic)*0.001
        elif unit=='m':
            ic=float(ic)*1000.0
        elif unit=='u':
            ic=float(ic)
        else:
            raise(exceptions)
        fie=float(re.search(r'USER    \(1\) Final Intermolecular Energy     = (?P<pp>[ 0-9.\-+]+) kcal/mol',pose_content1).group('pp'))
        vhde=float(re.search(r'USER        vdW \+ Hbond \+ desolv Energy     = (?P<pp>[ 0-9.\-+]+) kcal/mol',pose_content1).group('pp'))
        ee=float(re.search(r'USER        Electrostatic Energy            = (?P<pp>[ 0-9.\-+]+) kcal/mol',pose_content1).group('pp'))
        ftie=float(re.search(r'USER    \(2\) Final Total Internal Energy     = (?P<pp>[ 0-9.\-+]+) kcal/mol',pose_content1).group('pp'))
        tfe=float(re.search(r'USER    \(3\) Torsional Free Energy           = (?P<pp>[ 0-9.\-+]+) kcal/mol',pose_content1).group('pp'))
        use=float(re.search(r'USER    \(4\) Unbound System\'s Energy         = (?P<pp>[ 0-9.\-+]+) kcal/mol',pose_content1).group('pp'))
        
        PDBHead='USER %d %f %f %f %f %f %f %f\n'%(
            CLUSTERRANK,feb,ic,fie,ee,ftie,tfe,use)
        DLGAtomLines=re.findall(r'ATOM .+?\n',pose_content1,re.S)
        for lineseq1 in range(len(PDBQTLines)):
            DLGAtomLines[lineseq1]=PDBQTLines[lineseq1][:30]+DLGAtomLines[lineseq1][30:54]+PDBQTLines[lineseq1][54:]
        PDBstr=PDBHead+''.join(DLGAtomLines)+'END\n'
        contentzib=zlib.compress(PDBstr)
        
        # Build Pose object
        pose1=POSE(
            rank=CLUSTERRANK,
            feb=feb,
            ic=ic,
            fie=fie,
            vhde=vhde,
            ee=ee,
            ftie=ftie,
            tfe=tfe,
            use=use,
            pdbfile=contentzib,
        )
        mol1.poses.append(pose1)
        posetag1=POSETAG(status=6)
        session.flush()
        pose1.tag=posetag1
        session.flush()
    session.close()

    